from __init__ import *
"""推理管道"""
import os
from typing import List, Dict, Any, Optional
from tqdm import tqdm

from core.image_processor import ImageProcessor
from dataloader.data_loader import DatasetLoader, LabelLoader
from utils.geometry import CoordinateConverter
from inference.detector_postprocess import DetectorPostProcessor
from evaluation.evaluator import DetectionEvaluator
from inference.base_engine import InferenceEngine

class InferencePipeline:
    """推理管道"""
    
    def __init__(self, 
                 engine: InferenceEngine,
                 postprocessor: DetectorPostProcessor,
                 evaluator: Optional[DetectionEvaluator] = None,
                 image_size: int = 640):
        self.engine = engine
        self.postprocessor = postprocessor
        self.evaluator = evaluator
        self.image_size = image_size
        self.image_processor = ImageProcessor()
    
    def run_single_image(self, 
                        img_path: str,
                        label_path: Optional[str] = None) -> Dict[str, Any]:
        """运行单张图像推理"""
        # 图像预处理
        img_tensor, img_orig, scale_factor, pad_w, pad_h = \
            self.image_processor.preprocess_image(img_path, self.image_size)
        
        # 推理
        raw_output = self.engine.infer(img_tensor)
        
        # 后处理
        detections = self.postprocessor.process_raw_output(
            raw_output, img_orig.shape, scale_factor, pad_w, pad_h
        )
        
        result = {
            'image_path': img_path,
            'detections': detections,
            'image_shape': img_orig.shape
        }
        
        # 如果有标签，添加评估信息
        if label_path and self.evaluator:
            labels = LabelLoader.load_label(label_path)
            ground_truths = CoordinateConverter.label_to_xyxy(labels, img_orig.shape)
            
            tp, fp, fn = self.evaluator.evaluate_single_image(detections, ground_truths)
            result.update({
                'ground_truths': ground_truths,
                'true_positives': tp,
                'false_positives': fp,
                'false_negatives': fn
            })
        
        return result
    
    def run_dataset(self, 
                   img_dir: str,
                   label_dir: Optional[str] = None,
                   img_ext: str = '*.jpg') -> Dict[str, Any]:
        """运行数据集推理"""
        # 加载数据集
        dataset_loader = DatasetLoader(img_dir, label_dir or img_dir, img_ext)
        image_paths = dataset_loader.get_image_paths()
        
        if not image_paths:
            raise ValueError(f"No images found in {img_dir}")
        
        # 重置评估器
        if self.evaluator:
            self.evaluator.reset()
        
        results = []
        
        # 批量推理
        for img_path in tqdm(image_paths, desc="Processing images"):
            label_path = None
            if label_dir:
                label_path = dataset_loader.get_label_path(img_path)
            
            result = self.run_single_image(img_path, label_path)
            results.append(result)
            
            # 添加到评估器
            if self.evaluator and label_path:
                self.evaluator.add_batch(
                    [result['detections']],
                    [result['ground_truths']],
                    [img_path]
                )
        
        # 计算整体指标
        summary = {'total_images': len(results)}
        
        if self.evaluator:
            metrics = self.evaluator.get_metrics()
            summary.update(metrics)
            
            # 计算COCO mAP
            try:
                map_50, map_50_95 = self.evaluator.calculate_coco_map(
                    self.postprocessor.num_classes
                )
                summary.update({
                    'mAP@0.5': map_50,
                    'mAP@0.5:0.95': map_50_95
                })
            except Exception as e:
                print(f"Warning: COCO mAP calculation failed: {e}")
        
        return {
            'results': results,
            'summary': summary
        }